"""
损失函数定义
作者: magical857
日期: 2025-11-01
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class KoopmanLoss(nn.Module):
    """Koopman网络的组合损失函数"""
    
    def __init__(self, config):
        super(KoopmanLoss, self).__init__()
        
        self.w_reconstruction = config.WEIGHT_RECONSTRUCTION
        self.w_prediction = config.WEIGHT_PREDICTION
        self.w_linearity = config.WEIGHT_LINEARITY
        self.w_consistency = config.WEIGHT_CONSISTENCY
        
        self.mse = nn.MSELoss()
    
    def reconstruction_loss(self, states_pred, states_true):
        """重构损失"""
        return self.mse(states_pred, states_true)
    
    def prediction_loss(self, koopman_pred, koopman_true):
        """Koopman空间预测损失"""
        if koopman_pred is None or koopman_true is None:
            return torch.tensor(0.0, device=koopman_pred.device if koopman_pred is not None else 'cpu')
        return self.mse(koopman_pred, koopman_true)
    
    def linearity_loss(self, koopman_embedding, koopman_operator):
        """
        Koopman线性约束损失
        确保: psi_{t+1} ≈ K @ psi_t
        """
        if koopman_embedding.shape[1] < 2:
            return torch.tensor(0.0, device=koopman_embedding.device)
        
        batch_size, seq_len, koopman_dim = koopman_embedding.shape
        
        # 提取相邻时间步
        psi_t = koopman_embedding[:, :-1, :].reshape(-1, koopman_dim)
        psi_t_plus_1_true = koopman_embedding[:, 1:, :].reshape(-1, koopman_dim)
        
        # Koopman算子预测
        psi_t_plus_1_pred = koopman_operator(psi_t)
        
        return self.mse(psi_t_plus_1_pred, psi_t_plus_1_true)
    
    def consistency_loss(self, reconstruction, koopman_embedding, decoder):
        """
        一致性损失: 确保编码-解码的一致性
        decode(encode(x)) ≈ x
        """
        reconstruction_from_koopman = decoder(koopman_embedding)
        return self.mse(reconstruction_from_koopman, reconstruction)
    
    def forward(self, model_outputs, targets, model):
        """
        Args:
            model_outputs: (reconstruction, koopman_embedding, koopman_prediction)
            targets: 真实状态
            model: 模型对象（用于访问组件）
        
        Returns:
            total_loss, loss_dict
        """
        reconstruction, koopman_embedding, koopman_prediction = model_outputs
        
        # 1. 重构损失
        loss_recon = self.reconstruction_loss(reconstruction, targets)
        
        # 2. Koopman预测损失
        if koopman_prediction is not None and koopman_embedding.shape[1] > 1:
            koopman_true = koopman_embedding[:, 1:, :]
            loss_pred = self.prediction_loss(koopman_prediction, koopman_true)
        else:
            loss_pred = torch.tensor(0.0, device=reconstruction.device)
        
        # 3. 线性约束损失
        loss_linearity = self.linearity_loss(
            koopman_embedding,
            model.koopman_operator
        )
        
        # 4. 一致性损失
        loss_consistency = self.consistency_loss(
            reconstruction,
            koopman_embedding,
            model.autoencoder.decoder
        )
        
        # 总损失
        total_loss = (
            self.w_reconstruction * loss_recon +
            self.w_prediction * loss_pred +
            self.w_linearity * loss_linearity +
            self.w_consistency * loss_consistency
        )
        
        loss_dict = {
            'total': total_loss.item(),
            'reconstruction': loss_recon.item(),
            'prediction': loss_pred.item(),
            'linearity': loss_linearity.item(),
            'consistency': loss_consistency.item()
        }
        
        return total_loss, loss_dict


if __name__ == "__main__":
    # 测试损失函数
    import sys
    sys.path.append('..')
    from config import Config
    from models.koopman_network import KoopmanNetwork
    
    # 创建模型
    model = KoopmanNetwork(
        state_dim=Config.STATE_DIM,
        control_dim=Config.CONTROL_DIM,
        hidden_dim=Config.LSTM_HIDDEN_DIM,
        num_layers=Config.LSTM_NUM_LAYERS,
        koopman_dim=Config.KOOPMAN_DIM,
        dropout=Config.LSTM_DROPOUT
    )
    
    # 创建损失函数
    criterion = KoopmanLoss(Config)
    
    # 模拟数据
    batch_size = 8
    seq_len = 50
    states = torch.randn(batch_size, seq_len, Config.STATE_DIM)
    controls = torch.randn(batch_size, seq_len, Config.CONTROL_DIM)
    
    # 前向传播
    reconstruction, koopman_emb, koopman_pred = model(states, controls)
    model_outputs = (reconstruction, koopman_emb, koopman_pred)
    
    # 计算损失
    total_loss, loss_dict = criterion(model_outputs, states, model)
    
    print("损失值:")
    for key, value in loss_dict.items():
        print(f"  {key}: {value:.6f}")